import torch
import torch.nn as nn


def my_conv_test():
    img_tensor = torch.ones((1, 3, 5, 5))
    img_tensor[:, 1, :, :] = img_tensor[:, 1, :, :] * 2
    img_tensor[:, 2, :, :] = img_tensor[:, 2, :, :] * 3

    conv_layer_1 = nn.Conv2d(3, 2, 3, bias=False, stride=1, padding=0,)
    conv_layer_1.weight.data = torch.ones(conv_layer_1.weight.shape)
    conv_layer_1.weight.data[1] = conv_layer_1.weight.data[1] * 2

    conv_layer_2 = nn.Conv2d(3, 2, (3, 3), padding=1, bias=False)
    conv_layer_2.weight.data = torch.ones(conv_layer_2.weight.shape)
    conv_layer_2.weight.data[1] = conv_layer_2.weight.data[1] * 2

    outputs_1 = conv_layer_1(img_tensor)
    # print(conv_layer_1.weight.data)
    print(outputs_1, outputs_1.shape)
    print("卷积前尺寸：{}\n卷积后尺寸：{}".format(img_tensor.shape, outputs_1.shape))

    outputs_2 = conv_layer_2(img_tensor)
    print(outputs_2, outputs_2.shape)
    print("卷积前尺寸：{}\n卷积后尺寸：{}".format(img_tensor.shape, outputs_2.shape))

if __name__ == "__main__":
    my_conv_test()
